import os
import paramiko
from paramiko import SSHClient
from getpass import getpass
import subprocess

# 获取服务器信息
server_address = input("请输入服务器IP或域名: ")
server_user = input("请输入服务器用户名: ")
server_port = input("请输入服务器的SSH端口（默认22）: ") or "22"
password = getpass("请输入服务器密码: ")

ssh_dir = os.path.expanduser("~/.ssh")
private_key_file = os.path.join(ssh_dir, "id_rsa")
public_key_file = os.path.join(ssh_dir, "id_rsa.pub")

# 1. 生成 SSH 密钥对（如果尚未生成）
if not os.path.exists(private_key_file) or not os.path.exists(public_key_file):
    print("正在生成SSH密钥...")
    subprocess.run(["ssh-keygen", "-t", "rsa", "-b", "4096", "-f", private_key_file, "-N", ""], check=True)
    print("密钥生成完毕")
else:
    print("SSH密钥已存在，跳过生成步骤")

# 2. 使用 paramiko 登录并配置免密登录
try:
    # 创建 SSH 客户端实例
    ssh = SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    # 使用密码连接到服务器
    ssh.connect(server_address, username=server_user, password=password, port=int(server_port))
    print("成功连接到服务器")

    # 创建 .ssh 目录
    stdin, stdout, stderr = ssh.exec_command('mkdir -p ~/.ssh && chmod 700 ~/.ssh')
    stdout.channel.recv_exit_status()  # 确保命令执行完成

    # 上传公钥
    sftp = ssh.open_sftp()
    sftp.put(public_key_file, '.ssh/temp_pub_key')
    sftp.close()

    # 添加公钥到 authorized_keys 并删除临时文件
    ssh.exec_command('cat ~/.ssh/temp_pub_key >> ~/.ssh/authorized_keys && rm ~/.ssh/temp_pub_key')
    ssh.exec_command('chmod 600 ~/.ssh/authorized_keys')
    print("公钥已添加到服务器并配置")

    # 测试免密登录
    print("正在测试免密登录...")
    ssh.close()  # 断开当前有密码的连接

    # 尝试免密登录
    key = paramiko.RSAKey.from_private_key_file(private_key_file)
    ssh.connect(server_address, username=server_user, pkey=key, port=int(server_port))
    print("免密登录成功！")

except Exception as e:
    print(f"发生错误: {e}")

finally:
    ssh.close()

# 3. 可选：设置 SSH 客户端配置文件
config_file = os.path.join(ssh_dir, "config")
if not os.path.exists(config_file):
    open(config_file, "w").close()

with open(config_file, "a") as f:
    f.write(f"\nHost {server_address}\n    HostName {server_address}\n    User {server_user}\n    Port {server_port}\n    IdentityFile {private_key_file}\n")

print(f"SSH配置文件已更新，可以通过 'ssh {server_address}' 进行快速登录")
